#include "libmri.h"
#include "auxfns.h"
#include <omp.h>

#include <iostream>
using std::cout;
using std::endl;
using std::cerr;

void
pocs2d_l1spirit_omp (cplxf *data, int *dsize, cplxf *k2d,
                     l1spirit_config const& config)
{
	int *ksize = config.ksize;
	int n_iters = config.l1spirit_iterations;
	float lambda_1 = config.l1_threshold_start,
	      lambda_n = config.l1_threshold_end;

	int xs0 = dsize[0], xs1 = dsize[1], xs2 = dsize[2], nch = dsize[3];
	int ks1 = ksize[1], ks2 = ksize[2];

	cplxf *X = new cplxf [xs1 * xs2 * nch];
	cplxf *K = new cplxf [xs1 * xs2 * nch * nch];

	fftwf_plan X_forward, X_inverse;
	fftwf_plan K_forward, K_inverse;

	plan_fft2c (X_forward, X_inverse, X, xs1, xs2, nch);
	plan_fft2c (K_forward, K_inverse, K, xs1, xs2, nch*nch);

	/* A little kludgey -- planning the wavelet needs to compute the
	 * calibration region size in order to determine the number of levels
	 * of wavelet decomposition to perform */

	float *lor, *hir, *lod, *hid;
	int nc, lf, L, *S;

	for (int i1 = 0; i1 < xs1*xs2*nch; ++i1)
		X [i1] = data [(xs0/2) + xs0*i1];

	plan_wavelet (nc, lf, L, S, lor, hir, lod, hid, X, xs1, xs2);
	l1spirit_debug (config,wavplan,0,S,2,2,L+2);

	float *mask;
	solver_mask (mask, data, xs0, xs1, xs2);

	l1spirit_debug (config,mask,0,mask,2,xs1,xs2);


	delete [] X;
	delete [] K;

	#pragma omp parallel
	{
		cplxf *X = new cplxf [xs1 * xs2 * nch];
		cplxf *Y = new cplxf [xs1 * xs2 * nch];
		cplxf *K = new cplxf [xs1 * xs2 * nch * nch];  

		complex *vmap = 0; 
		float *emap = 0;
		int qcols = config.calib_pixel_n_eig;
		int qriters = config.calib_pixel_n_qrit;

		bool do_orthiter = (config.calib_algorithm == l1spirit_config::cheev)
		                or (config.calib_algorithm == l1spirit_config::lanczos)
		                or (config.calib_algorithm == l1spirit_config::culanczos);

		if (do_orthiter){
			vmap = new complex [xs1*xs2*nch*qcols];
			emap = new float [xs1*xs2*qcols];
		}

		int nwa = S[0]*S[1];
		cplxf *C = new cplxf [nc * nch];
		float *absC = new float [nc], *TabsC = new float[nc];


		#pragma omp for
		for (int i0 = 0; i0 < xs0; ++i0)
		{
			//double t_i0 = timestamp();

			for (int i1 = 0; i1 < xs1*xs2*nch; ++i1)
				X [i1] = data [i0 + xs0*i1];

			if (i0 == config.debug_slice)
				l1spirit_debug(config,l1s_x0,i0,X,3,xs1,xs2,nch);

			for (int i = 0; i < xs1*xs2*nch*nch; ++i)
				K[i] = 0.0f;

			int k1_0 = pad_index (xs1, ks1),
			    k2_0 = pad_index (xs2, ks2);

			for (int ch = 0; ch < nch*nch; ++ch)
				for (int k2 = 0; k2 < ks2; ++k2)
					for (int k1 = 0; k1 < ks1; ++k1)
						K [k1_0+k1 + xs1*(k2_0+k2 + xs2*ch)]
						  = k2d [i0 + xs0*(ks1-1-k1 + ks1*(ks2-1-k2 + ks2*ch))]
						     * sqrt (xs1*xs2);

			fft2c_exec (K_inverse, 1.0f/sqrt(xs1*xs2), K, xs1, xs2, nch*nch);

			if (do_orthiter)
			{
				/*
				#pragma omp critical
				{
					cout << "[" << i0 << "]" << " calib2d_orthiter()" << endl;
				}
				*/
				double ts = timestamp();

				int imsize[3] = {xs1, xs2, nch};
				calib2d_orthiter (vmap, emap, imsize, (complex*) K, qcols, qriters);


			/*
				if (i0 == config.debug_slice){
					debug_shm ("K", K, 4, xs1, xs2, nch, nch);
					debug_shm ("vmap", vmap, 4, xs1, xs2, nch, qcols);
					debug_shm ("emap", emap, 3, xs1, xs2, qcols);
				}
			*/

				complex *G = new complex [nch*nch];
				complex *V = new complex [nch*qcols];
				complex _1[1] = {{1.0f,0.0f}}, _0[1] = {{0.0f,0.0f}};;

				for (int iy = 0; iy < xs2; ++iy)
					for (int ix = 0; ix < xs1; ++ix)
					{
						int j_1 = -1;
						float diff = HUGE_VALF;

						for (int j = 0; j < qcols; ++j){
							float d = fabsf (emap [ix + xs1*(iy + xs2*j)]-1.0f);
							if (d < diff){
								j_1 = j;
								diff = d;
							}
						}

						for (int i = 0; i < nch; ++i)
							V[i] = vmap [ix + xs1*(iy + xs2*i)];

						cgemm ('N', 'C', nch, nch, 1,
						        _1, V,nch, V,nch, 
						        _0, G,nch);

						for (int i = 0; i < nch*nch; ++i)
							K [ix + xs1*(iy + xs2*i)] = G [i];
					}

				delete [] G;
				delete [] V;

				ts = timestamp()-ts;

				if (config.verbose_info){
					#pragma omp critical
					cout << "[" << omp_get_thread_num() << ", " << i0 << "] "
					     << "orthiter: " << ts << " s" << endl;
				}
			}

			if (i0 == config.debug_slice)
				l1spirit_debug(config,l1s_kern,i0,K,4,xs1,xs2,nch,nch);


		/****************
		 * Projections Over Convex Sets (POCS) Iterations. We have 3 sets:
		 *  #1: Calibration-Consistent Images
		 *  #2: Wavelet-Sparse Images
		 *  #3: Data-consistent Images
		 */
			double ts_pocs = timestamp();
			for (int iter = 1; iter <= n_iters; ++iter)
			{
			/*************
			 * Projection onto Set #1 (Calibration Consistency). We implement
			 * this projection via a Power Method -- assuming that the SPIRiT
			 * G operator only has eigenvalues of 1, then the matrix-vector
			 * multiplication G*x projects onto G's eigenvectors. For
			 * efficiency, we perform the SPIRiT interpolation in the
			 * image-domain. */

				fft2c_exec (X_inverse, 1.0f/sqrt(xs1*xs2), X, xs1, xs2, nch);

				if (i0 == config.debug_slice)
					l1spirit_debug(config,l1s_ifft,iter,X,3,xs1,xs2,nch);

				spirit (Y, false, K, xs1, xs2, nch, nch, X, xs1, xs2, nch);

				if (i0 == config.debug_slice)
					l1spirit_debug(config,l1s_spir,iter,Y,3,xs1,xs2,nch);

			/************
			 * Projection onto Set #2 (Wavelet Sparsity). We implement this
			 * via the soft-thresholding operation, and we threshold the
			 * cross-coil magnitude of each pixel. We do not threshold the
			 * lowest-order Wavelet approximation coefficients (of which there
			 * are nwa).
			 */
				float lambda = lambda_1 + (iter-1.0f)*(lambda_n-lambda_1)
				                                      /(n_iters-1.0f);
				if (i0 == config.debug_slice)
					l1spirit_debug(config,l1s_lambda,iter,&lambda,2,1,1);

				if (lambda > 0)
				{
					for (int ch = 0; ch < nch; ++ch)
						forward_wavelet (&C [ch*nc], nc,
						                 &Y [ch*xs1*xs2], xs1, xs2,
						                 L, S, lod, hid, lf);

					if (i0 == config.debug_slice)
						l1spirit_debug(config,l1s_fwav,iter,C,3,1,nc,nch);

					for (int i = nwa; i < nc; ++i)
					{
						absC[i] = 0.0f;

						for (int ch = 0; ch < nch; ++ch)
							absC[i] += abs2 (C[i + nc*ch]);
						absC[i] = sqrt (absC [i]);

						if (0.0f > (TabsC[i] = absC[i] - lambda))
							TabsC[i] = 0.0f;

						for (int ch = 0; ch < nch; ++ch)
							C[i + nc*ch] *= TabsC[i] / (absC[i] + 1e-6f);
					}

					if (i0 == config.debug_slice)
						l1spirit_debug(config,l1s_thr,iter,C,3,1,nc,nch);

					for (int ch = 0; ch < nch; ++ch)
						inverse_wavelet (&X [ch*xs1*xs2], xs1, xs2,
						                 &C [ch*nc], nc,
						                 L, S, lor, hir, lf);

					if (i0 == config.debug_slice)
						l1spirit_debug(config,l1s_iwav,iter,X,3,xs1,xs2,nch);
				}
				else {
					memcpy (X, Y, xs1*xs2*nch*sizeof(*X));
				}

			/************
			 * Projection onto Set #3 (Data Consistency). This operation 
			 * forces the k-space values in the solution to be equal to the
			 * acquired k-space values, and leaves other k-space locations alone
			 */
				fft2c_exec (X_forward, 1.0f/sqrt(xs1*xs2), X, xs1, xs2, nch);

				if (i0 == config.debug_slice)
					l1spirit_debug(config,l1s_fft,iter,X,3,xs1,xs2,nch);

				for (int ch = 0; ch < nch; ++ch)
					for (int i2 = 0; i2 < xs2; ++i2)
						for (int i1 = 0; i1 < xs1; ++i1)
						{
							X [i1 + xs1*(i2 + xs2*ch)]
							      = mask[i1 + xs1*i2]
							         * data [i0 + xs0*(i1 + xs1*(i2 + xs2*ch))]
							      + (1.0f-mask[i1 + xs1*i2])
							         * X [i1 + xs1*(i2 + xs2*ch)];
						}

				if (i0 == config.debug_slice)
					l1spirit_debug(config,l1s_proj,iter,X,3,xs1,xs2,nch);
			}

			for (int i1 = 0; i1 < xs1*xs2*nch; ++i1)
				data [i0 + xs0*i1] = X [i1];

			ts_pocs = timestamp()-ts_pocs;
			if (config.verbose_info){
				#pragma omp critical
					cout << "[" << omp_get_thread_num() << ", " << i0 << "] "
					     << "pocsiter: " << ts_pocs << " s" << endl;
			}
		}

		delete [] X;
		delete [] Y;
		delete [] K;
		delete [] C;

		if (do_orthiter){
			delete [] vmap;
			delete [] emap;
		}
		
	}
}


